import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from algorithm import Generator
from algorithm import Discriminator
from algorithm2 import Generator,Discriminator



# 定义生成器和判别器模型
generator = Generator()
discriminator = Discriminator()

# 加载保存的模型状态字典
generator.load_state_dict(torch.load('./generator-2.pth'))
discriminator.load_state_dict(torch.load('./discriminator-2.pth'))

test_points = torch.tensor(np.random.uniform(5, 45, (100, 2)), dtype=torch.float32)  # 生成一些新的位置
generated_rss = generator(test_points)
min_rss = -80
max_rss = -20
generated_rss = min_rss + (max_rss - min_rss) / 2 * (generated_rss + 1)
print("Generated RSS values for new points:")
print(test_points)
print(generated_rss)

# 可视化生成的RSS值
test_points = test_points.detach().numpy()
generated_rss = generated_rss.detach().numpy()

plt.scatter(test_points[:, 0], test_points[:, 1], c=np.linalg.norm(generated_rss, axis=1), cmap='viridis')
plt.colorbar(label="Generated RSS (Norm of RSS values)")
plt.title("Generated RSS values for new points")
plt.xlabel("X Position")
plt.ylabel("Y Position")
plt.show()
